Accelerator model state dict #7474

shuyingsunshine21 · 2021-05-10T19:13:03Z

What does this PR do?

Before submitting

Was this discussed/approved via a GitHub issue? (not for typos and docs)
Did you read the contributor guideline, Pull Request section?
Did you make sure your PR does only one thing, instead of bundling different changes together?
Did you make sure to update the documentation with your changes? (if necessary)
Did you write any new necessary tests? (not for typos and docs)
Did you verify new and existing tests pass locally with your changes?
Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

Is this pull request ready for review? (if not, please submit in draft mode)
Check that all items from Before submitting are resolved
Make sure the title is self-explanatory and the description concisely explains the PR
Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

…lightning pull latest code

…oint_consolidate Update test_all_gather_grad.py

This reverts commit 9d4a2b8.

…1-checkpoint_consolidate" This reverts commit c5053da, reversing changes made to 0d23d75.

This reverts commit 0d23d75.

This reverts commit 70fe5da.

This reverts commit a9aae99.

This reverts commit ea74906.

This reverts commit bf70e43.

This reverts commit f172101.

This reverts commit 536c132.

This reverts commit 3a9fde9.

This reverts commit 7a369f4.

…lightning

This reverts commit 8222dc9.

This reverts commit 6c095b2.

This reverts commit 250d0aa.

This reverts commit 8651d54.

This reverts commit dcdcd29.

…-lightning

…lightning

codecov · 2021-05-10T19:14:21Z

Codecov Report

Merging #7474 (adec445) into master (f6fe715) will decrease coverage by 4%.
The diff coverage is 100%.

@@           Coverage Diff           @@
##           master   #7474    +/-   ##
=======================================
- Coverage      92%     88%    -4%     
=======================================
  Files         199     199            
  Lines       13030   13057    +27     
=======================================
- Hits        11989   11431   -558     
- Misses       1041    1626   +585

for more information, see https://pre-commit.ci

SeanNaren

cc @awaelchli @tchaton

pytorch_lightning/plugins/training_type/training_type_plugin.py

ananthsub

do you see a good way to unit test the state dict being exposed from the accelerator?

ananthsub · 2021-05-11T04:11:45Z

pytorch_lightning/accelerators/accelerator.py

@@ -420,6 +420,12 @@ def optimizer_state(self, optimizer: Optimizer) -> Dict[str, Tensor]:
 """
 return getattr(self.training_type_plugin, 'optimizer_state', lambda x: x.state_dict())(optimizer)

+ def state_dict(self) -> Dict[str, Union[Any, Tensor]]:


why the union with tensor? same for L429

Suggested change

def state_dict(self) -> Dict[str, Union[Any, Tensor]]:

def state_dict(self) -> Dict[str, Any]:

Oh, I just follow the same type defined below for on_save, but I think you are right, Dict[str, Any] should be enough

@ananthsub I added the Union here originally, I think. This was done since it often is a tensor, but doesn't have to. And if you omit the tensor here, IDEs don't show the autocomplete-suggestions for Tensors

pytorch_lightning/plugins/training_type/training_type_plugin.py

shuyingsunshine21 · 2021-05-11T07:06:13Z

do you see a good way to unit test the state dict being exposed from the accelerator?

will check if we already have some existing test for state dict, if not, i will come up with something.

pytorch_lightning/plugins/training_type/training_type_plugin.py

justusschock · 2021-05-11T10:07:37Z

pytorch_lightning/accelerators/accelerator.py

@@ -420,6 +420,12 @@ def optimizer_state(self, optimizer: Optimizer) -> Dict[str, Tensor]:
 """
 return getattr(self.training_type_plugin, 'optimizer_state', lambda x: x.state_dict())(optimizer)

+ def state_dict(self) -> Dict[str, Union[Any, Tensor]]:


@ananthsub I added the Union here originally, I think. This was done since it often is a tensor, but doesn't have to. And if you omit the tensor here, IDEs don't show the autocomplete-suggestions for Tensors

Shuying Sun and others added 30 commits March 23, 2021 12:06

Fix some test errors

89f284d

Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags:

Merge branch 'master' of https://github.com/PyTorchLightning/pytorch-…

80cfbff

…lightning pull latest code

checkpoint consolidation

536c132

Update ddp_spawn.py

f172101

Update test_metric_result_integration.py

bf70e43

Update test_results.py

ea74906

Update utils.py

a9aae99

Update utils.py

70fe5da

Update test_all_gather_grad.py

0d23d75

Update test_all_gather_grad.py

ca6f98b

Merge pull request #1 from shuyingsunshine21/shuyingsunshine21-checkp…

c5053da

…oint_consolidate Update test_all_gather_grad.py

Update test_results.py

9d4a2b8

Revert "Update test_results.py"

7635b4f

This reverts commit 9d4a2b8.

Revert "Merge pull request #1 from shuyingsunshine21/shuyingsunshine2…

d64f90c

…1-checkpoint_consolidate" This reverts commit c5053da, reversing changes made to 0d23d75.

Revert "Update test_all_gather_grad.py"

dcdcd29

This reverts commit 0d23d75.

Revert "Update utils.py"

8651d54

This reverts commit 70fe5da.

Revert "Update utils.py"

15f4b9e

This reverts commit a9aae99.

Revert "Update test_results.py"

250d0aa

This reverts commit ea74906.

Revert "Update test_metric_result_integration.py"

6c095b2

This reverts commit bf70e43.

Revert "Update ddp_spawn.py"

8222dc9

This reverts commit f172101.

Revert "checkpoint consolidation"

3a9fde9

This reverts commit 536c132.

Revert "Revert "checkpoint consolidation""

7a369f4

This reverts commit 3a9fde9.

Revert "Revert "Revert "checkpoint consolidation"""

b4a0b9e

This reverts commit 7a369f4.

Merge branch 'master' of https://github.com/PyTorchLightning/pytorch-…

5cf1db1

…lightning

Revert "Revert "Update ddp_spawn.py""

0ce7e05

This reverts commit 8222dc9.

Revert "Revert "Update test_metric_result_integration.py""

fe9736d

This reverts commit 6c095b2.

Revert "Revert "Update test_results.py""

c314ef6

This reverts commit 250d0aa.

Revert "Revert "Update utils.py""

c3feda0

This reverts commit 8651d54.

Revert "Revert "Update test_all_gather_grad.py""

c759477

This reverts commit dcdcd29.

Merge branch 'master' of https://github.com/shuyingsunshine21/pytorch…

7a8e540

…-lightning

Shuying Sun added 5 commits April 13, 2021 21:21

Merge branch 'master' of https://github.com/PyTorchLightning/pytorch-…

f337156

…lightning

Merge branch 'master' of https://github.com/PyTorchLightning/pytorch-…

35bb931

…lightning

modify model state dict to training type plugin

efa39ff

rebase and pre-commit

47e2269

remove changes

d79c971

shuyingsunshine21 requested review from awaelchli, Borda, carmocca, justusschock, kaushikb11, SeanNaren, tchaton and williamFalcon as code owners May 10, 2021 19:13

Shuying Sun and others added 3 commits May 10, 2021 12:16

add changelog

497c002

fixing isort for pre-commit failure

465e967

[pre-commit.ci] auto fixes from pre-commit.com hooks

d5dd739

for more information, see https://pre-commit.ci

SeanNaren approved these changes May 10, 2021

View reviewed changes

awaelchli added distributed Generic distributed-related topic feature Is an improvement or enhancement labels May 10, 2021

awaelchli added this to the v1.4 milestone May 10, 2021

shuyingsunshine21 commented May 11, 2021

View reviewed changes

pytorch_lightning/plugins/training_type/training_type_plugin.py Outdated Show resolved Hide resolved

ananthsub reviewed May 11, 2021

View reviewed changes

awaelchli approved these changes May 11, 2021

View reviewed changes

pytorch_lightning/plugins/training_type/training_type_plugin.py Outdated Show resolved Hide resolved

justusschock approved these changes May 11, 2021

View reviewed changes

Address code review

adec445

ananthsub approved these changes May 11, 2021

View reviewed changes

SeanNaren merged commit 8538c1f into Lightning-AI:master May 11, 2021

shuyingsunshine21 deleted the accelerator_model_state_dict branch May 11, 2021 18:22

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Accelerator model state dict #7474

Accelerator model state dict #7474

shuyingsunshine21 commented May 10, 2021 •

edited

Loading

codecov bot commented May 10, 2021 •

edited

Loading

SeanNaren left a comment

ananthsub left a comment

ananthsub May 11, 2021

shuyingsunshine21 May 11, 2021

justusschock May 11, 2021

shuyingsunshine21 commented May 11, 2021

justusschock May 11, 2021

	def state_dict(self) -> Dict[str, Union[Any, Tensor]]:
	def state_dict(self) -> Dict[str, Any]:

Accelerator model state dict #7474

Accelerator model state dict #7474

Conversation

shuyingsunshine21 commented May 10, 2021 • edited Loading

What does this PR do?

Before submitting

PR review

Did you have fun?

codecov bot commented May 10, 2021 • edited Loading

Codecov Report

SeanNaren left a comment

Choose a reason for hiding this comment

ananthsub left a comment

Choose a reason for hiding this comment

ananthsub May 11, 2021

Choose a reason for hiding this comment

shuyingsunshine21 May 11, 2021

Choose a reason for hiding this comment

justusschock May 11, 2021

Choose a reason for hiding this comment

shuyingsunshine21 commented May 11, 2021

justusschock May 11, 2021

Choose a reason for hiding this comment

shuyingsunshine21 commented May 10, 2021 •

edited

Loading

codecov bot commented May 10, 2021 •

edited

Loading